import numpy as np
from sklearn.tree import DecisionTreeClassifier, export_text
import re

# 西瓜数据集（特征 + 标签）
dataSet = [
    ['青绿', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', '好瓜'],
    ['乌黑', '蜷缩', '沉闷', '清晰', '凹陷', '硬滑', '好瓜'],
    ['乌黑', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', '好瓜'],
    ['青绿', '蜷缩', '沉闷', '清晰', '凹陷', '硬滑', '好瓜'],
    ['浅白', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', '好瓜'],
    ['青绿', '稍蜷', '浊响', '清晰', '稍凹', '软粘', '好瓜'],
    ['乌黑', '稍蜷', '浊响', '稍糊', '稍凹', '软粘', '好瓜'],
    ['乌黑', '稍蜷', '浊响', '清晰', '稍凹', '硬滑', '好瓜'],
    ['乌黑', '稍蜷', '沉闷', '稍糊', '稍凹', '硬滑', '坏瓜'],
    ['青绿', '硬挺', '清脆', '清晰', '平坦', '软粘', '坏瓜'],
    ['浅白', '硬挺', '清脆', '模糊', '平坦', '硬滑', '坏瓜'],
    ['浅白', '蜷缩', '浊响', '模糊', '平坦', '软粘', '坏瓜'],
    ['青绿', '稍蜷', '浊响', '稍糊', '凹陷', '硬滑', '坏瓜'],
    ['浅白', '稍蜷', '沉闷', '稍糊', '凹陷', '硬滑', '坏瓜'],
    ['乌黑', '稍蜷', '浊响', '清晰', '稍凹', '软粘', '坏瓜'],
    ['浅白', '蜷缩', '浊响', '模糊', '平坦', '硬滑', '坏瓜'],
    ['青绿', '蜷缩', '沉闷', '稍糊', '稍凹', '硬滑', '坏瓜'],
]

# 特征映射定义
feature_mapping = {
    '色泽': {'青绿': 0, '乌黑': 1, '浅白': 2},
    '根蒂': {'蜷缩': 0, '稍蜷': 1, '硬挺': 2},
    '敲声': {'浊响': 0, '沉闷': 1, '清脆': 2},
    '纹理': {'清晰': 0, '稍糊': 1, '模糊': 2},
    '脐部': {'凹陷': 0, '稍凹': 1, '平坦': 2},
    '触感': {'硬滑': 0, '软粘': 1}
}

# 特征值反向映射（用于决策树可视化）
value_mapping = {
    '色泽': {0: '青绿', 1: '乌黑', 2: '浅白'},
    '根蒂': {0: '蜷缩', 1: '稍蜷', 2: '硬挺'},
    '敲声': {0: '浊响', 1: '沉闷', 2: '清脆'},
    '纹理': {0: '清晰', 1: '稍糊', 2: '模糊'},
    '脐部': {0: '凹陷', 1: '稍凹', 2: '平坦'},
    '触感': {0: '硬滑', 1: '软粘'}
}

# 标签映射
label_mapping = {'好瓜': 1, '坏瓜': 0}

# 特征名称（用于决策树可视化）
feature_names = list(feature_mapping.keys())

# 转换数据集为数值形式
X = []
y = []

for sample in dataSet:
    # 转换特征
    features = []
    for i, feature_value in enumerate(sample[:-1]):
        feature_name = feature_names[i]
        features.append(feature_mapping[feature_name][feature_value])

    # 转换标签
    label = label_mapping[sample[-1]]

    X.append(features)
    y.append(label)

# 转换为NumPy数组
X = np.array(X)
y = np.array(y)

# 创建并训练决策树模型
clf = DecisionTreeClassifier(max_depth=4, random_state=42)
clf.fit(X, y)

# 获取决策树文本表示
tree_text = export_text(clf, feature_names=feature_names)


# 将决策树文本中的数值条件转换为中文描述
def convert_conditions(tree_text):
    # 替换分类结果
    tree_text = tree_text.replace("class: 0", "分类结果: 坏瓜")
    tree_text = tree_text.replace("class: 1", "分类结果: 好瓜")

    # 替换特征条件
    for feature in feature_names:
        # 匹配模式：特征名 <= 或 > 数值
        pattern = r"({})\s*(<=|>)\s*([0-9.]+)".format(feature)

        def replace_match(match):
            feat_name = match.group(1)
            operator = match.group(2)
            value = float(match.group(3))

            # 获取最接近的整数值
            int_value = round(value)
            if int_value in value_mapping[feat_name]:
                return "{} {} {}".format(
                    feat_name,
                    operator,
                    value_mapping[feat_name][int_value]
                )
            return match.group(0)

        tree_text = re.sub(pattern, replace_match, tree_text)

    # 添加中文标题和格式优化
    formatted_text = "### 决策树结构\n\n"
    lines = tree_text.split("\n")

    for line in lines:
        # 替换缩进符号
        line = line.replace("|---", "├──")
        line = line.replace("|   ", "│   ")
        line = line.replace("|", "│")

        # 添加层级标题
        if "├──" in line:
            indent_level = line.count("│   ") // 2
            if indent_level == 0:
                line = line.replace("├──", "#### 分支 ")
            elif indent_level == 1:
                line = line.replace("├──", "├── 分支 ")
            else:
                line = line.replace("├──", "│   ├── 分支 ")

        formatted_text += line + "\n"

    return formatted_text


# 打印美化后的决策树
formatted_tree = convert_conditions(tree_text)
print(formatted_tree)

# 测试样本预测
test_samples = [
    ['青绿', '蜷缩', '浊响', '清晰', '凹陷', '硬滑'],  # 好瓜
    ['乌黑', '稍蜷', '沉闷', '稍糊', '稍凹', '软粘'],  # 坏瓜
    ['浅白', '硬挺', '清脆', '模糊', '平坦', '硬滑'],  # 坏瓜
]

print("\n测试样本预测结果:")
for i, sample in enumerate(test_samples, 1):
    # 转换测试样本
    test_features = []
    for j, feature_value in enumerate(sample):
        feature_name = feature_names[j]
        test_features.append(feature_mapping[feature_name][feature_value])

    test_features = np.array([test_features])

    # 预测并显示结果
    prediction = clf.predict(test_features)
    print(f"样本{i}: {sample} => {'好瓜' if prediction[0] == 1 else '坏瓜'}")